'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2020-06-24 17:44:13
'''
from subprocess import *
import random, time, json, sys

class DateRange:
    def __init__(self, bgn, end):
        import datetime
        bgn = bgn.split('-')
        begin = datetime.date(int(bgn[0]), int(bgn[1]), int(bgn[2]))
        end = end.split('-')
        end = datetime.date(int(end[0]), int(end[1]), int(end[2]))
        d = begin
        delta = datetime.timedelta(days=1)
        self.dates = []
        while d <= end:
            self.dates.append(d.strftime("%Y-%m-%d"))
            d += delta

class BatchTrainJson():
    def __init__(self, model_name, ps_url, process_num, dates, f_path, net_name, labels):
        self.model_name = model_name
        self.batch_size = 1000
        self.process_num = process_num
        self.shuffle_size = 20000
        self.watch_num = 50000
        self.neg_rate = 0.99
        self.ps_url = ps_url
        self.run_shell = '/data1/xuejian/sync/galaxy/galaxy'
        self.dates = dates
        self.f_path = f_path
        self.random_sleep = 10
        self.net_conf = net_name
        self.labels = labels
        self.shell_p = self.run_shell + " -layer /data1/xuejian/sync/model_zoo/models/galaxy/" + self.net_conf + " -mn " + self.model_name + " -ps_url " + self.ps_url + ':' + str(sys.argv[1]) + " -w " + str(self.watch_num) + " -s " + str(self.shuffle_size) + " -b " + str(self.batch_size) + " -ne " + str(self.neg_rate) + " -labels " + self.labels + " -is_json true"
        print('shell = ', self.shell_p)
        
        


    def train(self):
        for date in self.dates:
            print(date)
            proc_dict = {}
            num_list = [0] * self.process_num
            print("num_list = ", num_list) 
            for i in range(self.process_num):
                if self.process_num <= 1:
                    f_date = self.f_path + date
                else:
                    f_date = self.f_path + date + '/' + str(i)
                if i != 0:
                    time.sleep(random.randint(0, self.random_sleep))
                proc = Popen('cat ' + f_date + ' | ' + self.shell_p, shell=True)
                print('process start')
                proc_dict[i] = proc
            for _, p in proc_dict.items():
                p.wait()

    def train2(self):
        
        proc = Popen(self.shell_p, shell=True, stdin=PIPE)
        print('process start')
        for date in self.dates:
            print(date)
            with open(self.f_path + date, 'rb') as f_date:
                print(self.f_path + date)
                line = f_date.readline()
                while line != b'':
                    if line == b'':
                        break
                    proc.stdin.write(line)
                    proc.stdin.flush()
                    line = f_date.readline()
        
        proc.stdin.close()
        proc.wait()





class BatchTrain():
    def __init__(self, model_name, ps_url, process_num, dates, f_path, net_name, labels):
        self.model_name = model_name
        self.batch_size = 1000
        self.process_num = process_num
        self.shuffle_size = 20000
        self.watch_num = 50000
        self.neg_rate = 0.99
        self.ps_url = ps_url
        self.run_shell = '/data1/xuejian/sync/galaxy/galaxy'
        self.dates = dates
        self.f_path = f_path
        self.random_sleep = 10
        self.net_conf = net_name
        self.labels = labels
        self.shell_p = self.run_shell + " -layer /data1/xuejian/sync/model_zoo/models/galaxy/" + self.net_conf + " -mn " + self.model_name + " -ps_url " + self.ps_url + ':' + str(sys.argv[1]) + " -w " + str(self.watch_num) + " -s " + str(self.shuffle_size) + " -b " + str(self.batch_size) + " -ne " + str(self.neg_rate) + " -labels " + self.labels
        print('shell = ', self.shell_p)
        
        


    def train(self):
        for date in self.dates:
            print(date)
            proc_dict = {}
            num_list = [0] * self.process_num
            print("num_list = ", num_list) 
            for i in range(self.process_num):
                if self.process_num <= 1:
                    f_date = self.f_path + date
                else:
                    f_date = self.f_path + date + '/' + str(i)
                if i != 0:
                    time.sleep(random.randint(0, self.random_sleep))
                proc = Popen('cat ' + f_date + ' | ' + self.shell_p, shell=True)
                print('process start')
                proc_dict[i] = proc
            for _, p in proc_dict.items():
                p.wait()

    def train2(self):
        
        proc = Popen(self.shell_p, shell=True, stdin=PIPE)
        print('process start')
        for date in self.dates:
            print(date)
            with open(self.f_path + date, 'rb') as f_date:
                print(self.f_path + date)
                line = f_date.readline()
                while line != b'':
                    if line == b'':
                        break
                    proc.stdin.write(line)
                    proc.stdin.flush()
                    line = f_date.readline()
        
        proc.stdin.close()
        proc.wait()

